'''
Created on Dec 6, 2009

@author: mkiyer
'''
from optparse import OptionParser
import glob
import logging
import os
import re
import sys
import pickle
import subprocess
import collections
from bx.intervals.intersection import Interval, IntervalTree
from veggie.gene.gene import parse_refseq_genes

def build_interval_trees(genes):
    gene_trees = collections.defaultdict(lambda: IntervalTree())    
    for g in genes:
        # add gene to gene tree
        gene_trees[g.chrom].insert_interval(Interval(g.tx_start, g.tx_end, value=g))
    return gene_trees

def build_tss_trees(genes, window=2000):
    gene_trees = collections.defaultdict(lambda: IntervalTree())    
    for g in genes:
        tss = g.tx_start if g.strand == '+' else g.tx_end            
        # add gene to gene tree
        gene_trees[g.chrom].insert_interval(Interval(tss - window, tss + window, value=g))
    return gene_trees

def fetch_gene_tss_mappings(genes):
    #tss_trees = collections.defaultdict(lambda: IntervalTree())
    tss_to_gene = collections.defaultdict(lambda: collections.defaultdict(lambda: set()))
    gene_to_tss = collections.defaultdict(lambda: set())
    for g in genes:
        tss = g.tx_start if g.strand == '+' else g.tx_end
        tss_to_gene[g.chrom][tss].add(g.symbol)
        gene_to_tss[g.symbol].add(tss)
    return tss_to_gene, gene_to_tss
    

def build_peak_trees(peaks, name, trees=None):
    if peaks is None:
        return
    if trees is None:
        trees = collections.defaultdict(lambda: IntervalTree())
    for p in peaks:
        chrom, start, end = p
        trees[chrom].insert_interval(Interval(start, end, value=name))
    return trees

def parse_peaks(peakfile):
    if peakfile is None:
        return
    for line in open(peakfile):
        if line is None:
            continue
        line = line.strip()
        if line is None:
            continue
        if line.startswith("track"):
            continue
        fields = line.split('\t')
        chrom, start, end = fields[0], int(fields[1]), int(fields[2])
        if start < 0:
            start = 0
        yield chrom, start, end

if __name__ == '__main__':
    logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s',
                        level=logging.DEBUG)
    logger = logging.getLogger(__name__)    
    # parse command line arguments
    optionparser = OptionParser("usage: %prog [options]")
    optionparser.add_option("--h3k4", dest="h3k4", default=None)
    optionparser.add_option("--h3k9", dest="h3k9", default=None)
    optionparser.add_option("--h3k27", dest="h3k27", default=None)
    optionparser.add_option("--h3k36", dest="h3k36", default=None)
    optionparser.add_option("-w", "--window", dest="window", type="int", default=2000)
    (options, args) = optionparser.parse_args()

    # build tree data structure with peaks
    peak_trees = collections.defaultdict(lambda: IntervalTree())    
    if options.h3k4 is not None:
        peak_trees = build_peak_trees(parse_peaks(options.h3k4), "H3K4me3", trees=peak_trees)
    if options.h3k9 is not None:
        peak_trees = build_peak_trees(parse_peaks(options.h3k9), "H3K9me3", trees=peak_trees)
    if options.h3k27 is not None:
        peak_trees = build_peak_trees(parse_peaks(options.h3k27), "H3K27me3", trees=peak_trees)
    if options.h3k36 is not None:
        peak_trees = build_peak_trees(parse_peaks(options.h3k36), "H3K36me3", trees=peak_trees)

    # match refseq genes with peaks
    genes = list(parse_refseq_genes())
    tss_size = options.window
    #tss2gene, gene2tss = fetch_gene_tss_mappings(genes)
    logging.debug("Searching reference genes for peaks...")
    for g in genes:
        tss = g.tx_start if g.strand == '+' else g.tx_end
        tss_window = (tss - tss_size, tss + tss_size)
        # intersect genes with peaks
        hits = peak_trees[g.chrom].find(tss_window[0], tss_window[1])
        histone_marks = set([hit.value for hit in hits])
        if len(histone_marks) == 0:
            continue
        output_str = [g.symbol,
                      g.acc,
                      '%s:%d-%d' % (g.chrom, tss_window[0], tss_window[1]),
                      ','.join(histone_marks)]
        sys.stdout.write('\t'.join(output_str) + '\n')
